from ultralytics import YOLO
import cv2
import torch
import numpy as np

# 加载模型
model = YOLO("yolo11n-seg.pt")   # 你的自定义分割模型

# 预测图片
results = model("./ultralytics/assets/zidane.jpg")

for result in results:
    # 原图
    img = result.orig_img.copy()
    h, w = img.shape[:2]  # 原图尺寸

    if result.masks is not None:
        masks = result.masks.data  # (num_objects, h_pred, w_pred)
        names = model.names
        boxes = result.boxes.xyxy
        cls_ids = result.boxes.cls

        for i in range(masks.shape[0]):
            # 掩膜 -> numpy
            mask = masks[i].cpu().numpy()

            # 将掩膜 resize 到原图大小
            mask_resized = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)

            # 生成随机颜色
            color = np.random.randint(0, 255, (3,), dtype=np.uint8)

            # 创建彩色掩膜
            colored_mask = np.zeros_like(img, dtype=np.uint8)
            colored_mask[mask_resized > 0.5] = color

            # 叠加掩膜到原图
            img = cv2.addWeighted(img, 1, colored_mask, 0.5, 0)

            # 绘制边框和标签
            x1, y1, x2, y2 = map(int, boxes[i])
            label = f"{names[int(cls_ids[i])]}"
            cv2.rectangle(img, (x1, y1), (x2, y2), color.tolist(), 2)
            cv2.putText(img, label, (x1, y1 - 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, color.tolist(), 2)

        # 保存可视化结果
        cv2.imwrite("segmentation_result.jpg", img)
        print("可视化结果已保存到 segmentation_result.jpg")